package com.shanlitech.et;

import java.util.List;
import java.util.ArrayList;
import java.util.Map;
import java.util.HashMap;
import java.util.Iterator;

import com.google.protobuf.Descriptors;
import com.google.protobuf.DynamicMessage;
import com.google.protobuf.AbstractMessage;
import com.google.protobuf.UninitializedMessageException;
import com.google.protobuf.InvalidProtocolBufferException;
import org.zeromq.ZMQ;
import org.zeromq.ZMsg;
import org.zeromq.ZFrame;
import org.zeromq.ZMQ.Socket;

public class ZProtobuf {
	private Map<String,Descriptors.Descriptor>		descriptors_map;

	public final static class RecvStruct {
		public ArrayList<byte[]>		envelopes;
		public boolean					separator;
		public DynamicMessage			msg;

		RecvStruct() {
			envelopes = null;
			separator = false;
			msg = null;
		}
	}

	public ZProtobuf(Descriptors.FileDescriptor fd) {
		descriptors_map = new HashMap<String,Descriptors.Descriptor>();
		for(Descriptors.Descriptor md : fd.getMessageTypes() ) {
			descriptors_map.put(md.getFullName(),md);
		}
	}
	public ZProtobuf(List<Descriptors.FileDescriptor> fdl) {
		descriptors_map = new HashMap<String,Descriptors.Descriptor>();
		for(Descriptors.FileDescriptor fd : fdl) {
			for(Descriptors.Descriptor md : fd.getMessageTypes() ) {
				descriptors_map.put(md.getFullName(),md);
			}
		}
	}

	public DynamicMessage buildMessageFromName(String name,byte[] data) {
		Descriptors.Descriptor md = descriptors_map.get(name);
		if( md != null ) {
			try {
				return DynamicMessage.parseFrom(md,data);
			} catch( InvalidProtocolBufferException e ) {
				System.out.println(e.getMessage());
			}
		} else {
			System.out.println("can not found Descriptor for " + name);
		}
		return null;
	}

	public boolean send(Socket sock,AbstractMessage msg,ArrayList<byte[]> envelopes,boolean separator) {
		try {

			ZMsg zm = new ZMsg();

			if( envelopes != null ) {
				separator = true;
				for(byte[] envelope : envelopes) {
					if( envelope == null ) return false;
					if( envelope.length == 0 ) return false;
					zm.add(new ZFrame(envelope));
				}
			}
			if( separator ) {
				zm.add(new ZFrame(""));
			}

			byte[] data = msg.toByteArray();
			zm.add(new ZFrame(msg.getDescriptorForType().getFullName()));
			zm.add(new ZFrame(data));

			return zm.send(sock);
		} catch( UninitializedMessageException e ) {
			return false;
		}
	}

	public boolean send(Socket sock,AbstractMessage msg,byte[] envelope) {
		if( envelope == null ) return false;
		if( envelope.length == 0 ) return false;

		ArrayList<byte[]> envelopes = new ArrayList<byte[]>();
		envelopes.add(envelope);
		return send(sock,msg,envelopes,true);
	}

	public boolean send(Socket sock,AbstractMessage msg,ArrayList<byte[]> envelopes) {
		return send(sock,msg,envelopes,true);
	}

	public boolean send(Socket sock,AbstractMessage msg,boolean separator) {
		return send(sock,msg,null,separator);
	}

	public boolean send(Socket sock,AbstractMessage msg) {
		return send(sock,msg,null,false);
	}


	public RecvStruct recv(Socket sock,boolean wait) {
		try {
			do {
				RecvStruct result = new RecvStruct();
				ZMsg zm = ZMsg.recvMsg(sock,wait);
				if( zm == null ) break;
				
				for(Iterator<ZFrame> it = zm.iterator(); it.hasNext(); ) {
					ZFrame f = it.next();
					if( f.hasData() && f.size() == 0 ) {
						result.separator = true;
						break;
					}
				}

				if( result.separator ) {
					result.envelopes = new ArrayList<byte[]>();
					do {
						ZFrame f = zm.pop();
						if( f.hasData() && f.size() == 0 ) {
							break;
						}
						result.envelopes.add(f.getData());
					} while(true);
				}

				ZFrame f = zm.pop();
				if( f == null ) break;

				String msg_name = f.getString(ZMQ.CHARSET);
				if( msg_name == null ) break;
				if( msg_name.length() == 0 ) break;

				f = zm.pop();
				if( f == null ) break;
				result.msg = buildMessageFromName(msg_name,f.getData());

				return result;
			} while(false);
		} catch ( UninitializedMessageException e ) {
		//} catch ( InvalidProtocolBufferException e ) {
		}
		return null;
	}
}
